package com.huawei.admins.platform.gateway.core.component.filter.factory;

import com.huawei.admins.platform.gateway.core.component.filter.FilterDefinition;
import com.huawei.admins.platform.gateway.core.component.filter.GatewayFilter;
import com.huawei.admins.platform.gateway.core.component.filter.GatewayFilterChain;
import com.huawei.admins.platform.gateway.core.config.GatewayProperties;
import com.huawei.admins.platform.gateway.core.util.ServerWebExchangeUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.reactivestreams.Publisher;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.HttpStatusCode;
import org.springframework.util.Assert;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;
import reactor.netty.Connection;
import reactor.retry.Backoff;
import reactor.retry.Repeat;
import reactor.retry.RepeatContext;
import reactor.retry.Retry;
import reactor.retry.RetryContext;

import java.io.IOException;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.TimeoutException;
import java.util.function.Predicate;
import java.util.function.Supplier;

public class RetryGatewayFilterFactory extends AbstractGatewayFilterFactory<RetryGatewayFilterFactory.RetryConfig> {

    public static final String RETRY_ITERATION_KEY = "retry_iteration";

    private static final Log log = LogFactory.getLog(RetryGatewayFilterFactory.class);

    public RetryGatewayFilterFactory(GatewayProperties properties) {
        super(properties);
    }

    @Override
    public GatewayFilter apply(String routeId, FilterDefinition definition) {

        RetryConfig retryConfig = loadConfiguration(routeId);
        retryConfig.validate();

        Repeat<ServerWebExchange> statusCodeRepeat = null;
        if (!retryConfig.getStatuses().isEmpty() || !retryConfig.getSeries().isEmpty()) {
            Predicate<RepeatContext<ServerWebExchange>> repeatPredicate = context -> {
                ServerWebExchange exchange = context.applicationContext();
                if (exceedsMaxIterations(exchange, retryConfig)) {
                    return false;
                }

                HttpStatusCode statusCode = exchange.getResponse().getStatusCode();

                boolean retryableStatusCode = retryConfig.getStatuses().contains(statusCode);

                // null status code might mean a network exception?
                if (!retryableStatusCode) {
                    // try the series
                    for (int i = 0; i < retryConfig.getSeries().size(); i++) {
                        if (statusCode instanceof HttpStatus httpStatus) {
                            if (httpStatus.series().equals(retryConfig.getSeries().get(i))) {
                                retryableStatusCode = true;
                                break;
                            }
                        }
                    }
                }

                final boolean finalRetryableStatusCode = retryableStatusCode;
                trace("retryableStatusCode: %b, statusCode %s, configured statuses %s, configured series %s",
                        () -> finalRetryableStatusCode, () -> statusCode, retryConfig::getStatuses,
                        retryConfig::getSeries);

                HttpMethod httpMethod = exchange.getRequest().getMethod();
                boolean retryableMethod = retryConfig.getMethods().contains(httpMethod);

                trace("retryableMethod: %b, httpMethod %s, configured methods %s", () -> retryableMethod,
                        () -> httpMethod, retryConfig::getMethods);
                return retryableMethod && finalRetryableStatusCode;
            };

            statusCodeRepeat = Repeat.onlyIf(repeatPredicate)
                    .doOnRepeat(context -> reset(context.applicationContext()));

            BackoffConfig backoff = retryConfig.getBackoff();
            if (backoff != null) {
                statusCodeRepeat = statusCodeRepeat.backoff(getBackoff(backoff));
            }
        }

        Retry<ServerWebExchange> exceptionRetry = null;
        if (!retryConfig.getExceptions().isEmpty()) {
            Predicate<RetryContext<ServerWebExchange>> retryContextPredicate = context -> {

                ServerWebExchange exchange = context.applicationContext();

                if (exceedsMaxIterations(exchange, retryConfig)) {
                    return false;
                }

                Throwable exception = context.exception();
                for (Class<? extends Throwable> retryableClass : retryConfig.getExceptions()) {
                    if (retryableClass.isInstance(exception)
                            || (exception != null && retryableClass.isInstance(exception.getCause()))) {
                        trace("exception or its cause is retryable %s, configured exceptions %s",
                                () -> getExceptionNameWithCause(exception), retryConfig::getExceptions);

                        HttpMethod httpMethod = exchange.getRequest().getMethod();
                        boolean retryableMethod = retryConfig.getMethods().contains(httpMethod);
                        trace("retryableMethod: %b, httpMethod %s, configured methods %s", () -> retryableMethod,
                                () -> httpMethod, retryConfig::getMethods);
                        return retryableMethod;
                    }
                }
                trace("exception or its cause is not retryable %s, configured exceptions %s",
                        () -> getExceptionNameWithCause(exception), retryConfig::getExceptions);
                return false;
            };
            exceptionRetry = Retry.onlyIf(retryContextPredicate)
                    .doOnRetry(context -> reset(context.applicationContext())).retryMax(retryConfig.getRetries());
            BackoffConfig backoff = retryConfig.getBackoff();
            if (backoff != null) {
                exceptionRetry = exceptionRetry.backoff(getBackoff(backoff));
            }
        }

        GatewayFilter gatewayFilter = apply(routeId, statusCodeRepeat, exceptionRetry);

        return new GatewayFilter() {

            @Override
            public int getOrder() {
                return 0;
            }

            @Override
            public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
                return gatewayFilter.filter(exchange, chain);
            }

            @Override
            public String toString() {
                return RetryGatewayFilterFactory.class.getName() +
                        "routeId=" + routeId +
                        "retries=" + retryConfig.getRetries() +
                        "series=" + retryConfig.getSeries() +
                        "statuses=" + retryConfig.getStatuses() +
                        "methods=" + retryConfig.getMethods() +
                        "exceptions=" + retryConfig.getExceptions();
            }
        };
    }

    public void reset(ServerWebExchange exchange) {
        Connection conn = exchange.getAttribute(ServerWebExchangeUtils.CLIENT_RESPONSE_CONN_ATTR);
        if (conn != null) {
            trace("disposing response connection before next iteration");
            conn.dispose();
            exchange.getAttributes().remove(ServerWebExchangeUtils.CLIENT_RESPONSE_CONN_ATTR);
        }
        ServerWebExchangeUtils.reset(exchange);
    }

    public GatewayFilter apply(String routeId, Repeat<ServerWebExchange> repeat, Retry<ServerWebExchange> retry) {
        return (exchange, chain) -> {
            trace("Entering retry-filter: " + routeId);

            // chain.filter returns a Mono<Void>
            Publisher<Void> publisher = chain.filter(exchange)
                    // .log("retry-filter", Level.INFO)
                    .doOnSuccess(aVoid -> updateIteration(exchange)).doOnError(throwable -> updateIteration(exchange));

            if (retry != null) {
                // retryWhen returns a Mono<Void>
                // retry needs to go before repeat
                publisher = ((Mono<Void>) publisher)
                        .retryWhen(reactor.util.retry.Retry.withThrowable(retry.withApplicationContext(exchange)));
            }
            if (repeat != null) {
                // repeatWhen returns a Flux<Void>
                // so this needs to be last and the variable a Publisher<Void>
                publisher = ((Mono<Void>) publisher).repeatWhen(repeat.withApplicationContext(exchange));
            }

            return Mono.fromDirect(publisher);
        };
    }

    public boolean exceedsMaxIterations(ServerWebExchange exchange, RetryConfig retryConfig) {
        Integer iteration = exchange.getAttribute(RETRY_ITERATION_KEY);

        boolean exceeds = iteration != null && iteration >= retryConfig.getRetries();
        trace("exceedsMaxIterations %b, iteration %d, configured retries %d", () -> exceeds, () -> iteration,
                retryConfig::getRetries);
        return exceeds;
    }

    @SafeVarargs
    private void trace(String message, Supplier<Object>... argSuppliers) {
        if (log.isTraceEnabled()) {
            Object[] args = new Object[argSuppliers.length];
            int i = 0;
            for (Supplier<Object> a : argSuppliers) {
                args[i] = a.get();
                ++i;
            }
            log.trace(String.format(message, args));
        }
    }

    private Backoff getBackoff(BackoffConfig backoff) {
        return Backoff.exponential(backoff.firstBackoff, backoff.maxBackoff, backoff.factor,
                backoff.basedOnPreviousValue);
    }

    private String getExceptionNameWithCause(Throwable exception) {
        if (exception != null) {
            StringBuilder builder = new StringBuilder(exception.getClass().getName());
            Throwable cause = exception.getCause();
            if (cause != null) {
                builder.append("{cause=").append(cause.getClass().getName()).append("}");
            }
            return builder.toString();
        }
        else {
            return "null";
        }
    }

    private void updateIteration(ServerWebExchange exchange) {
        int iteration = exchange.getAttributeOrDefault(RETRY_ITERATION_KEY, -1);
        int newIteration = iteration + 1;
        trace("setting new iteration in attr %d", () -> newIteration);
        exchange.getAttributes().put(RETRY_ITERATION_KEY, newIteration);
    }

    private static <T> List<T> toList(T... items) {
        return new ArrayList<>(Arrays.asList(items));
    }

    @Override
    @SuppressWarnings("unchecked")
    protected RetryConfig buildConfig(String routeId, Map<String, String> filterArgs) {
        RetryConfig retryConfig = new RetryConfig();

        Optional.ofNullable(filterArgs.get("retries")).filter(NumberUtils::isCreatable).map(Integer::parseInt).ifPresent(retryConfig::setRetries);
        String retries = filterArgs.get("retries");
        if (NumberUtils.isCreatable(retries)) {
            retryConfig.setRetries(Integer.parseInt(retries));
        }

        List<HttpStatus.Series> series = new ArrayList<>();
        List<HttpStatus> statues = new ArrayList<>();
        List<HttpMethod> methods = new ArrayList<>();
        List<Class<? extends Throwable>> exceptions = new ArrayList<>();

        for (Map.Entry<String, String> argEntry : filterArgs.entrySet()) {
            String key = argEntry.getKey();
            String value = argEntry.getValue();
            if (StringUtils.isBlank(key) || StringUtils.isBlank(value)) {
                continue;
            }
            if (key.startsWith("series")) {
                series.add(HttpStatus.Series.valueOf(value));
            } else if (key.startsWith("statuses")) {
                statues.add(HttpStatus.valueOf(value));
            } else if (key.startsWith("methods")) {
                methods.add(HttpMethod.valueOf(value));
            } else if (key.startsWith("exceptions")) {
                try {
                    Class<?> aClass = Class.forName(value);
                    if (Throwable.class.isAssignableFrom(aClass)) {
                        exceptions.add((Class<? extends Throwable>) aClass);
                    }
                } catch (ClassNotFoundException ignored) {
                }
            }
        }
        return retryConfig
                .setSeries(series.toArray(new HttpStatus.Series[0]))
                .setStatuses(statues.toArray(new HttpStatus[0]))
                .setMethods(methods.toArray(new HttpMethod[0]))
                .setExceptions(exceptions.toArray(new Class[0]));
    }

    public static class RetryConfig {

        private int retries = 3;

        private List<HttpStatus.Series> series = toList(HttpStatus.Series.SERVER_ERROR);

        private List<HttpStatus> statuses = new ArrayList<>();

        private List<HttpMethod> methods = toList(HttpMethod.GET);

        private List<Class<? extends Throwable>> exceptions = toList(IOException.class, TimeoutException.class);

        private BackoffConfig backoff;

        public RetryConfig allMethods() {
            return setMethods(HttpMethod.values());
        }

        public void validate() {
            Assert.isTrue(this.retries > 0, "retries must be greater than 0");
            Assert.isTrue(!this.series.isEmpty() || !this.statuses.isEmpty() || !this.exceptions.isEmpty(),
                    "series, status and exceptions may not all be empty");
            Assert.notEmpty(this.methods, "methods may not be empty");
            if (this.backoff != null) {
                this.backoff.validate();
            }
        }

        public BackoffConfig getBackoff() {
            return backoff;
        }

        public RetryConfig setBackoff(BackoffConfig backoff) {
            this.backoff = backoff;
            return this;
        }

        public RetryConfig setBackoff(Duration firstBackoff, Duration maxBackoff, int factor,
                                      boolean basedOnPreviousValue) {
            this.backoff = new BackoffConfig(firstBackoff, maxBackoff, factor, basedOnPreviousValue);
            return this;
        }

        public int getRetries() {
            return retries;
        }

        public RetryConfig setRetries(int retries) {
            this.retries = retries;
            return this;
        }

        public List<HttpStatus.Series> getSeries() {
            return series;
        }

        public RetryConfig setSeries(HttpStatus.Series... series) {
            this.series = Arrays.asList(series);
            return this;
        }

        public List<HttpStatus> getStatuses() {
            return statuses;
        }

        public RetryConfig setStatuses(HttpStatus... statuses) {
            this.statuses = Arrays.asList(statuses);
            return this;
        }

        public List<HttpMethod> getMethods() {
            return methods;
        }

        public RetryConfig setMethods(HttpMethod... methods) {
            this.methods = Arrays.asList(methods);
            return this;
        }

        public List<Class<? extends Throwable>> getExceptions() {
            return exceptions;
        }

        public RetryConfig setExceptions(Class<? extends Throwable>... exceptions) {
            this.exceptions = Arrays.asList(exceptions);
            return this;
        }

    }

    public static class BackoffConfig {

        private Duration firstBackoff = Duration.ofMillis(5);

        private Duration maxBackoff;

        private int factor = 2;

        private boolean basedOnPreviousValue = true;

        public BackoffConfig() {
        }

        public BackoffConfig(Duration firstBackoff, Duration maxBackoff, int factor, boolean basedOnPreviousValue) {
            this.firstBackoff = firstBackoff;
            this.maxBackoff = maxBackoff;
            this.factor = factor;
            this.basedOnPreviousValue = basedOnPreviousValue;
        }

        public void validate() {
            Assert.notNull(this.firstBackoff, "firstBackoff must be present");
        }

        public Duration getFirstBackoff() {
            return firstBackoff;
        }

        public void setFirstBackoff(Duration firstBackoff) {
            this.firstBackoff = firstBackoff;
        }

        public Duration getMaxBackoff() {
            return maxBackoff;
        }

        public void setMaxBackoff(Duration maxBackoff) {
            this.maxBackoff = maxBackoff;
        }

        public int getFactor() {
            return factor;
        }

        public void setFactor(int factor) {
            this.factor = factor;
        }

        public boolean isBasedOnPreviousValue() {
            return basedOnPreviousValue;
        }

        public void setBasedOnPreviousValue(boolean basedOnPreviousValue) {
            this.basedOnPreviousValue = basedOnPreviousValue;
        }

    }
}
